#ifndef SIMD_BITSTREAM_H_MODULE
#define SIMD_BITSTREAM_H_MODULE

/*-****************************************
*  Dependencies
******************************************/
#include "mem.h"            /* unaligned access routines */
#include "compiler.h"       /* UNLIKELY() */
#include "debug.h"          /* assert(), DEBUGLOG(), RAWLOG() */
#include "error_private.h"  /* error codes and messages */
#include "bits.h"           /* ZSTD_highbit32 */
#include "simd_x86.h"

typedef struct {
    SimdBitContainerType bitContainer;
    SimdSize bitsConsumed;
    SimdSize index;
    SimdSize limitIndex;
    const char* start;
} SIMD_BIT_DStream_t;


typedef enum { SIMD_BIT_DStream_unfinished = 0,  /* fully refilled */
               SIMD_BIT_DStream_endOfBuffer = 1, /* still some bits left in bitstream */
               SIMD_BIT_DStream_completed = 2,   /* bitstream entirely consumed, bit-exact */
               SIMD_BIT_DStream_overflow = 3     /* user requested more bits than present in bitstream */
    } SIMD_BIT_DStream_status;  /* result of BIT_reloadDStream() */

MEM_STATIC size_t   SIMD_BIT_initDStream(SIMD_BIT_DStream_t* bitD, const void* srcBuffer, size_t *offset);
MEM_STATIC SimdSize   SIMD_BIT_readBits(SIMD_BIT_DStream_t* bitD, SimdSize nbBits);
MEM_STATIC SIMD_BIT_DStream_status SIMD_BIT_reloadDStream(SIMD_BIT_DStream_t* bitD);
MEM_STATIC SimdSize SIMD_BIT_endOfDStream(const SIMD_BIT_DStream_t* bitD);

/*=====    Local Constants   =====*/
static const size_t SIMD_BIT_mask[] = {
    0,          1,         3,         7,         0xF,       0x1F,
    0x3F,       0x7F,      0xFF,      0x1FF,     0x3FF,     0x7FF,
    0xFFF,      0x1FFF,    0x3FFF,    0x7FFF,    0xFFFF,    0x1FFFF,
    0x3FFFF,    0x7FFFF,   0xFFFFF,   0x1FFFFF,  0x3FFFFF,  0x7FFFFF,
    0xFFFFFF,   0x1FFFFFF, 0x3FFFFFF, 0x7FFFFFF, 0xFFFFFFF, 0x1FFFFFFF,
    0x3FFFFFFF, 0x7FFFFFFF}; /* up to 31 bits */
#define SIMD_BIT_MASK_SIZE (sizeof(SIMD_BIT_mask) / sizeof(SIMD_BIT_mask[0]))



/*-********************************************************
*  bitStream decoding
**********************************************************/
/*! BIT_initDStream() :
 *  Initialize a BIT_DStream_t.
 * `bitD` : a pointer to an already allocated BIT_DStream_t structure.
 * `srcBuffer` must be the *exact* size of the bitStream, in bytes.
 * `offset` must be an array with PARALLEL_DEGREE+1 size
 * @return : size of stream (== srcSize), or an errorCode if a problem is detected
 */
MEM_STATIC size_t SIMD_BIT_initDStream(SIMD_BIT_DStream_t* bitD, const void* srcBuffer, size_t *offset)
{
    if (offset[PARALLEL_DEGREE] < 1 || offset[0] != 0) { ZSTD_memset(bitD, 0, sizeof(*bitD)); return ERROR(srcSize_wrong); }
    void* left_offset = (void*)offset;
    void* right_offset = (void*)(offset + 1);
    SIMD_LOADER_DEF(simd_loader);
    SimdSize container_size = SIMD_SET1_SIZE(sizeof(bitD->bitContainer));

    bitD->start = (const char*)srcBuffer;
    bitD->limitIndex = container_size;
    bitD->index = SIMD_LOAD_SIZE(right_offset, simd_loader);

    SimdSize bits_per_byte = SIMD_SET1_SIZE(8);
    SimdSize start_index = SIMD_LOAD_SIZE(left_offset, simd_loader);
    SimdSize length = SIMD_SUB_SIZE(bitD->index, start_index);
    bitD->index = SIMD_SUB_SIZE(SIMD_MAX_SIZE(bitD->index, SIMD_ADD_SIZE(start_index, container_size)), container_size);
    bitD->bitContainer = SIMD_GATHER_SIZE(srcBuffer, bitD->index);

    SimdSize end_mark_bits_num = SIMD_SET1_SIZE(0);
    const BYTE* byte_src = srcBuffer;
    for(int i = 1;i <= PARALLEL_DEGREE; i++) {
        BYTE last_byte = byte_src[offset[i]-1];
        int bits_num = last_byte? 8 - ZSTD_highbit32(last_byte):0;
        end_mark_bits_num = SIMD_INSERT_SIZE(end_mark_bits_num, bits_num, i);
        if(last_byte == 0) return ERROR(corruption_detected);
    }
    bitD->bitsConsumed = SIMD_ADD_SIZE(
        SIMD_MULLO_SIZE(SIMD_SUB_SIZE(container_size, SIMD_MIN_SIZE(container_size, length)), bits_per_byte),
        end_mark_bits_num
    );
    return 0;
}



FORCE_INLINE_TEMPLATE SimdSize SIMD_BIT_getMiddleBits(SimdBitContainerType bitContainer, SimdSize start, SimdSize nbBits)
{
    SimdSize regMask = SIMD_SET1_SIZE(sizeof(BitContainerLane)*8 - 1);
    SimdSize right_shift = SIMD_AND_SIZE(start, regMask);
    SimdSize bit_mask = SIMD_GATHER_SIZE((long long int const *)SIMD_BIT_mask, nbBits);
    return SIMD_AND_SIZE(SIMD_SRL_SIZE(bitContainer, right_shift), bit_mask);
}

/*! SIMD_BIT_lookBits() :
 *  Provides next n bits from local register.
 *  local register is not modified.
 *  On 32-bits, maxNbBits==24.
 *  On 64-bits, maxNbBits==56.
 * @return : value extracted */
FORCE_INLINE_TEMPLATE SimdSize SIMD_BIT_lookBits(const SIMD_BIT_DStream_t*  bitD, SimdSize nbBits)
{
    /* if bitD->bitsConsumed + nbBits > sizeof(bitD->bitContainer)*8,
     * bitstream is likely corrupted, and result is undefined */
    SimdSize length = SIMD_SET1_SIZE(sizeof(BitContainerLane) * 8);
    SimdSize start = SIMD_SUB_SIZE(SIMD_SUB_SIZE(length, bitD->bitsConsumed), nbBits);
    return SIMD_BIT_getMiddleBits(bitD->bitContainer, (sizeof(bitD->bitContainer)*8) - bitD->bitsConsumed - nbBits, nbBits);
}


/*! SIMD_BIT_lookBitsFast() :
 *  unsafe version; only works if nbBits >= 1 */
MEM_STATIC SimdSize SIMD_BIT_lookBitsFast(const SIMD_BIT_DStream_t* bitD, SimdSize nbBits)
{
    SimdSize regMask = SIMD_SET1_SIZE(sizeof(BitContainerLane)*8 - 1);
    //assert(nbBits >= 1);
    SimdSize left_shift = SIMD_AND_SIZE(bitD->bitsConsumed, regMask);
    SimdSize right_shift = SIMD_AND_SIZE(SIMD_SUB_SIZE(SIMD_SET1_SIZE(sizeof(BitContainerLane)*8), nbBits), regMask);
    return SIMD_SRL_SIZE(SIMD_SLL_SIZE(bitD->bitContainer, left_shift), right_shift);
}

FORCE_INLINE_TEMPLATE void SIMD_BIT_skipBits(SIMD_BIT_DStream_t* bitD, SimdSize nbBits)
{
    bitD->bitsConsumed = SIMD_ADD_SIZE(bitD->bitsConsumed, nbBits);
}

/*! SIMD_BIT_readBits() :
 *  Read (consume) next n bits from local register and update.
 *  Pay attention to not read more than nbBits contained into local register.
 * @return : extracted value. */
FORCE_INLINE_TEMPLATE SimdSize SIMD_BIT_readBits(SIMD_BIT_DStream_t* bitD, SimdSize nbBits)
{
    SimdSize value = SIMD_BIT_lookBits(bitD, nbBits);
    SIMD_BIT_skipBits(bitD, nbBits);
    return value;
}

/*! BIT_readBitsFast() :
 *  unsafe version; only works if nbBits >= 1 */
MEM_STATIC SimdSize SIMD_BIT_readBitsFast(SIMD_BIT_DStream_t* bitD, SimdSize nbBits)
{
    SimdSize value = SIMD_BIT_lookBitsFast(bitD, nbBits);
    //assert(nbBits >= 1);
    SIMD_BIT_skipBits(bitD, nbBits);
    return value;
}

/*! BIT_reloadDStream_internal() :
 *  Simple variant of BIT_reloadDStream(), with two conditions:
 *  1. bitstream is valid : bitsConsumed <= sizeof(bitD->bitContainer)*8
 *  2. look window is valid after shifted down : bitD->ptr >= bitD->start
 */
MEM_STATIC SIMD_BIT_DStream_status SIMD_BIT_reloadDStream_internal(SIMD_BIT_DStream_t* bitD)
{
    //assert(bitD->bitsConsumed <= sizeof(bitD->bitContainer)*8);
    bitD->index = SIMD_SUB_SIZE(bitD->index, SIMD_SRL_SIZE(bitD->bitsConsumed, SIMD_SET1_SIZE(3)));
    //assert(bitD->ptr >= bitD->start);
    bitD->bitsConsumed = SIMD_AND_SIZE(bitD->bitsConsumed, SIMD_SET1_SIZE(7));
    bitD->bitContainer = SIMD_GATHER_SIZE((long long int const *)bitD->start, bitD->index);
    return SIMD_BIT_DStream_unfinished;
}

/*! BIT_endOfDStream() :
 * @return : 1 if DStream has _exactly_ reached its end (all bits consumed).
 */
MEM_STATIC SimdSize SIMD_BIT_endOfDStream(const SIMD_BIT_DStream_t* DStream)
{
    return SIMD_AND_SIZE(
        SIMD_EQ_SIZE(DStream->index, SIMD_SET1_SIZE(0)), 
        SIMD_EQ_SIZE(DStream->bitsConsumed, SIMD_SET1_SIZE(sizeof(BitContainerLane)*8))
    );
}
#endif